import paddle
import random as np

class restBlock(paddle.nn.Layer):
    
    def __init__(self,in_channels=256,out_channels=256):
        super(restBlock,self).__init__()
        self.conv1 = paddle.nn.Conv2D(in_channels,out_channels,3,stride=1, padding=1)
        self.conv2 = paddle.nn.Conv2D(in_channels,out_channels,3,stride=1, padding=1)
        self.relu = paddle.nn.ReLU()
        
    def forward(self, x):
        y = self.conv1(x)
        y = self.relu(y)
        y = self.conv2(y)
        x = x*0.1
        y = paddle.add(x,y)
        return y


class baseline_layer(paddle.nn.Layer):
    
    def __init__(self,in_channels=3):
        super(baseline_layer,self).__init__()
        self.conv1 = paddle.nn.Conv2D(in_channels,256,3,stride=1, padding=1)
        self.conv2 = paddle.nn.Conv2D(256,256,3,stride=1, padding=1)
        for i in range(1,33):
            exec("self.rb_%d = restBlock(256,256)"%(i))
            
    def forward(self, x):
        y = self.conv1(x)
        x = y
        for i in range(1,33):
            exec("y=self.rb_%d(y)"%(i))
        y = self.conv2(y)
        y = paddle.add(x,y)
        return y


class upsampleX2(paddle.nn.Layer):
    
    def __init__(self,in_channels=256):
        super(upsampleX2,self).__init__()
        self.conv1 = paddle.nn.Conv2D(in_channels,12,3,stride=1, padding=1)
        self.conv2 = paddle.nn.Conv2D(3,3,3,stride=1, padding=1)
        
    def forward(self, x):
        y = self.conv1(x)
        y=paddle.nn.functional.pixel_shuffle(y,2)
        y = self.conv2(y)
        return y
  
    
class upsampleX3(paddle.nn.Layer):
    
    def __init__(self,in_channels=256):
        super(upsampleX3,self).__init__()
        self.conv1 = paddle.nn.Conv2D(in_channels,27,3,stride=1, padding=1)
        self.conv2 = paddle.nn.Conv2D(3,3,3,stride=1, padding=1)
        
    def forward(self, x):
        y = self.conv1(x)
        y = paddle.nn.functional.pixel_shuffle(y,3)
        y = self.conv2(y)
        return y
    
    
class upsampleX4(paddle.nn.Layer):
    
    def __init__(self,in_channels=256):
        super(upsampleX4,self).__init__()
        self.conv1 = paddle.nn.Conv2D(in_channels,48,3,stride=1, padding=1)
        self.conv2 = paddle.nn.Conv2D(12,12,3,stride=1, padding=1)
        self.conv3 = paddle.nn.Conv2D(3,3,3,stride=1, padding=1)
        
    def forward(self, x):
        y = self.conv1(x)
        y = paddle.nn.functional.pixel_shuffle(y,2)
        y = self.conv2(y)
        y = paddle.nn.functional.pixel_shuffle(y,2)
        y = self.conv3(y)
        return y


class EDSR(paddle.nn.Layer):
    
    def __init__(self,scale=2):
        super(EDSR,self).__init__()
        self.scale = scale
        self.base = baseline_layer(in_channels=3)
        self.up2 = upsampleX2(in_channels=256)
        self.up3 = upsampleX3(in_channels=256)
        self.up4 = upsampleX4(in_channels=256)
        
    def forward(self,x):
        y = self.base(x)
        if self.scale == 2:
            y = self.up2(y)
        elif self.scale == 3:
            y = self.up3(y)
        elif self.scale == 4:
            y = self.up4(y)
        else:
            raise ValueError("scale value only 2 or 3 or 4")
        return y
        

if __name__ == "__main__":
    #edsr = EDSR(2)
    #imgs = np.random.random((1,3,48,48))
    #p = edsr(imgs)
    pass




























    









